# Copyright 2020 - 2022 MONAI Consortium
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#     http://www.apache.org/licenses/LICENSE-2.0
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import torch
import torch.nn as nn

from .swin_uneter import SwinTransformer as SwinViT
from monai.utils import ensure_tuple_rep
from monai import transforms
from einops import rearrange

transform = transforms.Compose(
    [
        transforms.RandAdjustContrast(prob=0.7, gamma=[0.7, 1.5]),
        transforms.RandShiftIntensity(0.4, prob=0.4),
        transforms.RandGaussianNoise(prob=0.5, mean=0.05, std=0.05),
    ]
)


def aug_rand(samples):
    samples = transform(samples)
    return samples


class SSLHead(nn.Module):
    def __init__(self, args):
        super(SSLHead, self).__init__()

        patch_size = ensure_tuple_rep(2, args.spatial_dims)
        window_size = ensure_tuple_rep(7, args.spatial_dims)

        self.swinViT = SwinViT(
            in_chans=args.in_channels,
            embed_dim=args.feature_size,
            window_size=window_size,
            patch_size=patch_size,
            depths=[2, 2, 2, 2],
            num_heads=[3, 6, 12, 24],
            mlp_ratio=4.0,
            qkv_bias=True,
            drop_rate=0.0,
            attn_drop_rate=0.0,
            drop_path_rate=args.dropout_path_rate,
            norm_layer=torch.nn.LayerNorm,
            use_checkpoint=args.use_checkpoint,
            spatial_dims=args.spatial_dims,
        )

        self.swinViT_momentum = SwinViT(
            in_chans=args.in_channels,
            embed_dim=args.feature_size,
            window_size=window_size,
            patch_size=patch_size,
            depths=[2, 2, 2, 2],
            num_heads=[3, 6, 12, 24],
            mlp_ratio=4.0,
            qkv_bias=True,
            drop_rate=0.0,
            attn_drop_rate=0.0,
            drop_path_rate=args.dropout_path_rate,
            norm_layer=torch.nn.LayerNorm,
            use_checkpoint=args.use_checkpoint,
            spatial_dims=args.spatial_dims,
        )

        for param_b, param_m in zip(self.swinViT.parameters(), self.swinViT_momentum.parameters()):
            param_m.data.copy_(param_b.data)  # initialize
            param_m.requires_grad = False  # not update by gradient

        self.contrastive_pre = nn.Identity()

        self._build_projector_and_predictor_mlps(256, 4096)

        self.T = 0.1

    @torch.no_grad()
    def _update_momentum_encoder(self, m):
        """Momentum update of the momentum encoder"""
        for param_b, param_m in zip(self.swinViT.parameters(), self.swinViT_momentum.parameters()):
            param_m.data = param_m.data * m + param_b.data * (1. - m)

        for param_b, param_m in zip(self.projector.parameters(), self.momentum_projector.parameters()):
            param_m.data = param_m.data * m + param_b.data * (1. - m)

    def contrastive_loss(self, q, k):
        # normalize
        q = nn.functional.normalize(q, dim=1)
        k = nn.functional.normalize(k, dim=1)
        # gather all targets
        # k = concat_all_gather(k)
        # Einstein sum is more intuitive
        logits = torch.einsum('nc,mc->nm', [q, k]) / self.T
        N = logits.shape[0]  # batch size per GPU
        labels = (torch.arange(N, dtype=torch.long) + N * 0).cuda()
        return nn.CrossEntropyLoss()(logits, labels) * (2 * self.T) + 1e-9

    def _build_mlp(self, num_layers, input_dim, mlp_dim, output_dim, last_bn=True):
        mlp = []
        for l in range(num_layers):
            dim1 = input_dim if l == 0 else mlp_dim
            dim2 = output_dim if l == num_layers - 1 else mlp_dim

            mlp.append(nn.Linear(dim1, dim2, bias=False))

            if l < num_layers - 1:
                mlp.append(nn.BatchNorm1d(dim2))
                mlp.append(nn.ReLU(inplace=True))
            elif last_bn:
                # follow SimCLR's design: https://github.com/google-research/simclr/blob/master/model_util.py#L157
                # for simplicity, we further removed gamma in BN
                mlp.append(nn.BatchNorm1d(dim2, affine=False))

        return nn.Sequential(*mlp)

    def _build_projector_and_predictor_mlps(self, dim, mlp_dim):
        # self.projector = self._build_mlp(3, 768, mlp_dim, dim)
        # self.momentum_projector = self._build_mlp(3, 768, mlp_dim, dim)
        
        # # predictor
        # self.predictor = self._build_mlp(2, dim, mlp_dim, dim)
        
        self.projector = nn.Sequential(
            nn.Linear(768, 1024, bias=False),
            # nn.BatchNorm1d(1024),
            nn.ReLU(inplace=True),
            nn.Linear(1024, 1024, bias=False),
            # nn.BatchNorm1d(1024),
            nn.ReLU(inplace=True),
            nn.Linear(1024, 512, bias=False),
        )
        self.momentum_projector = nn.Sequential(
            nn.Linear(768, 1024, bias=False),
            # nn.BatchNorm1d(1024),
            nn.ReLU(inplace=True),
            nn.Linear(1024, 1024, bias=False),
            # nn.BatchNorm1d(1024),
            nn.ReLU(inplace=True),
            nn.Linear(1024, 512, bias=False),
        )
        
        # predictor
        self.predictor = nn.Sequential(
            nn.Linear(512, 1024, bias=False),
            # nn.BatchNorm1d(1024),
            nn.ReLU(inplace=True),
            nn.Linear(1024, 512, bias=False),
        )

    def forward(self, x, y):
        # [b, 768, 4, 4, 4]
        b, _, _, _, _ = x.size()
        
        # [BCHWZ] -> [BC]
        x = self.swinViT(x)[-1].flatten(2).permute(0, 2, 1).mean(1).squeeze(1)
        y = self.swinViT(y)[-1].flatten(2).permute(0, 2, 1).mean(1).squeeze(1)

        q1 = self.predictor(self.projector(x))
        q2 = self.predictor(self.projector(y))

        with torch.no_grad():  # no gradient
            self._update_momentum_encoder(0.999)  # update the momentum encoder

            # compute momentum features as targets
            k1 = self.momentum_projector(x)
            k2 = self.momentum_projector(y)

        return self.contrastive_loss(q1, k2) + self.contrastive_loss(q2, k1)


@torch.no_grad()
def concat_all_gather(tensor):
    """
    Performs all_gather operation on the provided tensors.
    *** Warning ***: torch.distributed.all_gather has no gradient.
    """
    tensors_gather = [torch.ones_like(tensor)
        for _ in range(torch.distributed.get_world_size())]
    torch.distributed.all_gather(tensors_gather, tensor, async_op=False)

    output = torch.cat(tensors_gather, dim=0)
    return output

    inps = torch.randn((1, 768, 4, 4, 4))
    
    print(inps.flatten(2).permute(0, 2, 1).mean(1).squeeze(1).size())
    
    # print(y.size())
    
    # projector = nn.Sequential(
    #     nn.Linear(768, 512),
    #     nn.BatchNorm1d(64),
    #     nn.ReLU(inplace=True),
    #     nn.Linear(512, 256),
    #     nn.BatchNorm1d(64),
    # )
    # print(projector(y).size())
    
    inps = torch.tensor([[1, 8],
                         [0, 0],
                         [-1, -8]], dtype=torch.float32)
    norm = nn.BatchNorm1d(2)
    
    print(inps)
    print(norm(inps))
    
    mean_ = torch.mean(inps, dim=0)
    var_ = torch.var(inps, dim=0, unbiased=False)
    print(mean_)
    print(var_)
    
    norm_ = torch.div(inps[0][1]-mean_[1], torch.sqrt(var_[1]))
    print(norm_)
